# websocket_manager.py
from fastapi import WebSocket
from typing import List, Dict


class ConnectionManager:
    def __init__(self):
        # 所有活跃连接
        self.active_connections: List[WebSocket] = []
        # 客户端ID -> 连接 的映射
        self.client_connections: Dict[str, WebSocket] = {}

    async def connect(self, websocket: WebSocket, client_id: str):
        """接受新连接"""
        await websocket.accept()
        self.active_connections.append(websocket)
        self.client_connections[client_id] = websocket
        print(f"客户端 {client_id} 已连接，当前连接数：{len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket, client_id: str):
        """移除断开连接"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        if client_id in self.client_connections:
            self.client_connections.pop(client_id)
        print(f"客户端 {client_id} 已断开，当前连接数：{len(self.active_connections)}")

    async def broadcast(self, message: str):
        """广播消息"""
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception as e:  # noqa: E722
                print(f"⚠️  广播消息失败: {str(e)}")

    async def send_to_client(self, client_id: str, message: str):
        """发消息给指定客户端"""
        if client_id in self.client_connections:
            try:
                await self.client_connections[client_id].send_text(message)
            except Exception as e:  # noqa: E722
                print(f"⚠️  发送消息给客户端 {client_id} 失败: {str(e)}")
                # 如果发送失败，移除该连接
                ws = self.client_connections.pop(client_id, None)
                if ws and ws in self.active_connections:
                    self.active_connections.remove(ws)
        else:
            print(f"⚠️  客户端 {client_id} 未连接或连接已断开")


wbs = ConnectionManager()